# email_spam.py
# 
# This file does the actual spam detection in Peter Ballard's email spam filter.
# The function check() takes in an email message (the "message" class from
# the standard Python module email), and returns a status
# ("wham", "mham", "spam" or "ham") as well as a one line description.
#
# Feel free to modify this program to your taste.
# But the intention is that you can use this program unchanged,
# and that you only need to modify email_defs.py
#
# Version 1.0, 4-Apr-2005.
#
#
#Copyright (C) 2005  Peter Ballard
#
#This program is free software; you can redistribute it and/or
#modify it under the terms of the GNU General Public License
#as published by the Free Software Foundation; either version 2
#of the License, or (at your option) any later version.
#
#This program is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with this program; if not, write to the Free Software
#Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
# 
import string
import re

import email_defs

##################################################
# Functions for getting data from another file
##################################################

# each (whitespace separated) string in the file becomes a list element,
# unless the line begins with "#" or ";"
def file_to_list(filename):
    list = []
    fp = open(filename, "r")
    for line in fp.readlines():
        words = string.split(line)
        if len(words) and words[0][0]!="#" and words[0][0]!=";":
            list.extend(words)
    fp.close()
    return list

##################################################
# Functions for separating out part of the message
##################################################

def getfield(msg, field):
    if msg.has_key(field):
        return string.strip(msg.get(field))
    else:
        return ""

def subject(msg):
    return getfield(msg, "Subject")

# body includes subject,
# and bodies of all sub-messages if it's multipart
def body(msg):
    # initial space + trailing CR to aid regexp matching
    retstring = " " + subject(msg) + "\n"
    if msg.is_multipart():
        for submsg in msg.get_payload():
            retstring = retstring + submsg.as_string()
    else:
        retstring = retstring + msg.get_payload()
    return retstring

def extract_addresses(astr):
    # return list of matching strings
    alist = []
    for obj in re.finditer("[-\w\.]+@[-\w\.]+", astr):
        addr = string.lower(obj.group())
        alist.append( addr)
        # this little bit increases time from about 65 to 75 secs for 5500 msgs
        # it's probably worth it to get more generalised addr checking
        parts = string.split(addr, "@")
        if len(parts)==2:
            # append all domain names,
            # down to domain names of two parts
            # e.g. with pballard@ozemail.com.au it would add "ozemail.com.au" and "com.au"
            bits = string.split(parts[1], ".")
            for i in range(0, len(bits)-1):
                alist.append(string.join( bits[i:], "."))
    return alist

# doing domains like this instead of regexps reduces time from 88 to 61 secs for 5764 messages
def extract_urls(astr):
    # return list of matching strings
    alist = []
    # checking for =2e doesnt cost much,
    # because some spams get identified quicker (before we reach regexps)
    for dot in [".", "=2e"]:
        for obj in re.finditer(r"http://[-\w" + dot + "]+", astr):
            addr = string.lower(obj.group())[7:] # strip 1st 7 chars i.e. "http://"
            # append all domain names,
            # down to domain names of one parts
            # e.g. http://ozemail.com.au it would add "ozemail.com.au" and "com.au" and "au"
            bits = string.split(addr, dot)
            for i in range(0, len(bits)):
                alist.append(string.join( bits[i:], "."))
    return alist

def recipient(msg):
    return extract_addresses(getfield(msg, "To")) + extract_addresses(getfield(msg, "Cc"))

def sender(msg):
    if msg.has_key("From"):
        # already in lower case coming out of extract_address[es]()
        alist = extract_addresses(msg.get("From"))
        if len(alist):
            return alist
    return []

# Content-Type string for all attachments, if any
def msg_attachment_types(msg):
    if msg.is_multipart():
        retlist = []
        for submsg in msg.get_payload():
            if submsg.has_key("Content-Type"):
                retlist.append(submsg.get("Content-Type"))
        return retlist
    else:
        return []

##################################################
# Functions for actual spam matching
##################################################

def file_to_phrase_list(filename):
    list = []
    fp = open(filename, "r")
    for line in fp.readlines():
        words = string.split(line)
        if len(words) and words[0][0]!="#" and words[0][0]!=";":
            list.append(string.strip(line))
    fp.close()
    return list

def really_from_me(msg):
    if find_listmember_in_list(email_defs.my_addresses, sender(msg)):
        xmailer = getfield(msg, "X-Mailer")
        if email_defs.really_from_me(xmailer):
            return email_defs.my_address
    else:
        return ""

def bad_attachment(msg, black_attachment_types):
    for strng in msg_attachment_types(msg):
        for ext in black_attachment_types:
            if re.search(r'\.' + ext, strng):
                return "##bad_attachment:." + ext
    return 0
    
def non_letter():
    return r"[^a-z0-9]"

def non_trickyletter():
    return r"[^a-z0-9@!\|]"

def domainchars():
    return r"[a-z0-9_\-]"

def nondomainchars():
    return r"[^a-z0-9_\-]"

def letteralt(ch):
    if ch=="a":
        # need to also add the accented a
        return "[a@]"
    elif ch=="i" or ch=="l" or ch=="1":
        # with ascii chars>127, it seems that raw strings dont work
        # need to also add the accented i
        # chr(195) = i with circumflex
        return "[il1!\\|" + chr(int("356",8)) + "]"
    if ch=="o":
        return "[o0]"
    else:
        return "[" + ch + "]"

def simpleregexp(str):
    return non_letter() + str + non_letter()

def trickyregexp(str):
    outstring = non_trickyletter() + letteralt(str[0])
    for i in range(1, len(str)):
        outstring = outstring + non_trickyletter() + "*?" + letteralt(str[i])
    return outstring + non_trickyletter()

def findword(str, wordlist):
    for word in wordlist:
        if re.search(simpleregexp(word), str):
            return simpleregexp(word)
    return 0

def findstring(str, wordlist):
    for word in wordlist:
        if string.count(str, word):
            return word
    return 0

def findregexp(str, wordlist):
    for word in wordlist:
        if re.search(word, str, re.MULTILINE): # makes no performance or result difference on current set
        #if re.search(word, str):
            return word
    return 0
    
def find_in_list(str, list):
    # using this rather list.count reduces time fro, 118 to 75 secs on 5586 messages
    if str in list:
        return str
    return ""

def find_listmember_in_list(strlist, list):
    for str in strlist:
        if str in list:
            return str
    return ""

def bad_recipients(recipient_list):
    #recipient_string = string.join(recipient_list, " ")
    if not find_listmember_in_list(email_defs.my_addresses, recipient_list):
        return "##recipient_not_me"
    elif recipient_list.count(email_defs.my_isp_domain) >= 3:
        return "##many_my_isp_recipients"
    else:
        return ""

def bad_subject(msg):
    if subject(msg)[:2]=="=?":
        return "=?"
    else:
        return ""

def empty_message(msg, body_mod):
    if subject(msg)=="" and string.strip(body_mod)=="":
        return "##empty_message"
    else:
        return ""

##################################################
# The main entry program
##################################################

def check(msg):    
    # strip any empty tags (i.e. <tag></tag>), and make lower case
    #body_notags = re.sub("<(?P<tag>[a-z]+)></(?P=tag)>", "", string.lower(body(msg)))

    bodylow = string.lower(body(msg))
    bodyearly = bodylow[:300] # first 200 chars in body (including subject)
    thissender = sender(msg)
    # look for whitelisted ham (wham)
    whitematch = (find_listmember_in_list(thissender, whitelist)
                  or really_from_me(msg))
    if whitematch:
        return "wham", whitematch
    # looked for "marked" ham (mham)
    mwhitematch = (findstring(bodyearly, whiteintrostrings)
                   or findregexp(bodyearly, whiteintroregexps))
    if mwhitematch:
        return "mham", mwhitematch
    # bad_recipients gets most hits so do it first
    blackmatch = (bad_recipients(recipient(msg))
                  or findstring(bodylow, blackstrings)
                  or empty_message(msg, bodylow)
                  or bad_attachment(msg, black_attachment_types)
                  or bad_subject(msg)
                  or findstring(extract_urls(bodylow), blackdomains)
                  # regexps are costliest so check them last (sped up 1099 emails from 43 to 24 secs!)
                  or findregexp(bodylow, blackregexps))
    if blackmatch:
        return "spam", blackmatch
    # do this last. It looks like a good trick but only caught 3 of 1508 spams
    body_notags = re.sub("<.*?>", "", bodylow)
    blackmatch = findregexp(body_notags, blackregexps)
    if blackmatch:
        return "spam", blackmatch
    else:
        return " ham", ""


##################################################
# Some top level definitions
##################################################
        
trickywords = email_defs.trickywords
blackwords = email_defs.blackwords
blackphrases = email_defs.blackphrases
blackregexps = email_defs.blackregexps

for word in trickywords:
    # this saves 1/10000, and simplifies some output, but again about 4/3 slowdown
    blackregexps.append(simpleregexp(word))
    blackregexps.append(trickyregexp(word))
    # this catches very few extra cases (in fact 1: "meds!"), but increases time 61 -> 82 secs for 6643 msgs
    #blackregexps.append(simpleregexp(word))

for word in blackwords:
    blackregexps.append(simpleregexp(word))

for blackphrase in blackphrases:
    list = string.split(string.strip(blackphrase))
    regexp = list[0]
    for i in range(1, len(list)):
        #regexp = regexp + r"\s" + list[i]
        regexp = regexp + non_letter() + "+?" + list[i]
    blackregexps.append(regexp)

blackstrings = email_defs.blackstrings

blackdomains = file_to_list(email_defs.blackdomains)

black_attachment_types = email_defs.black_attachment_types

# words in the intro, or early in the body, which indicate wham
whiteintrowords = email_defs.whiteintrowords

whiteintroregexps = []
for word in whiteintrowords:
    whiteintroregexps.append(simpleregexp(word))

# strings in the subject line which indicate mail is wham
whiteintrostrings = email_defs.whiteintrostrings

whitelist = file_to_list(email_defs.whitelist)
